#include "warray/warray_la.h"
#include "test.h"

#include <iostream>

#include "custom_type.h"
#include "helper.h"

typedef std::complex<float> cfloat;
typedef std::complex<double> cdouble;

typedef boost::mpl::list<int,float,double> scalar_types;
typedef boost::mpl::list<cfloat,cdouble> complex_types;
typedef boost::mpl::list<int,float,double,cfloat,cdouble> all_types;

BOOST_AUTO_TEST_CASE_TEMPLATE( s_matrixmult, T, all_types )
{
	T adata[9]     = {1,2,3,       4,5,6,       7,8,9};
	T bdata[9]     = {10,11,12,    13,14,15,    16,17,18};
	T ans[][9]     = {{84,90,96,    201,216,231, 318,342,366},  // a*b
		              {174,186,198, 213,228,243, 252,270,288},  // a'*b
		              {68,86,104,   167,212,257, 266,338,410},  // a*b'
		              {138,174,210, 171,216,261, 204,258,312}}; // a'*b'

	T ans1d[][3] = {{68, 167, 266},   // a*x,
		            {138, 171, 204}}; // a'*x

	//
	// Matrix-matrix multiplication
	//
	warp::arraynd<T, 2> A(warp::extents[3][3], adata);
	warp::arraynd<T, 2> B(warp::extents[3][3], bdata);
	warp::arraynd<T, 2> C(warp::extents[3][3]);
             
	T* c = C.c_array();

	// A*B = C
	warp::la::matrixmult(A, B, C);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[0],ans[0]+9);

	// A'*B = C
	warp::la::matrixmult(A, B, C, warp::blas::transpose, warp::blas::dont_transpose);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[1],ans[1]+9);

	// A*B' = C
	warp::la::matrixmult(A, B, C, warp::blas::dont_transpose, warp::blas::transpose);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[2],ans[2]+9);

	// A'*B' = C
	warp::la::matrixmult(A, B, C, warp::blas::transpose, warp::blas::transpose);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[3],ans[3]+9);

	//
	// Matrix-vector multiplication (square matrix)
	//
	{
		warp::arraynd<T, 1> x(warp::extents[3], bdata);
		warp::arraynd<T, 1> b(warp::extents[3]);
		c = b.c_array();

		// A*x = b
		warp::la::matrixmult(A, x, b);
		BOOST_CHECK_EQUAL_COLLECTIONS(c,c+3,ans1d[0],ans1d[0]+3);

		// A'*x = b
		warp::la::matrixmult(A, x, b, warp::blas::transpose);;
		BOOST_CHECK_EQUAL_COLLECTIONS(c,c+3,ans1d[1],ans1d[1]+3);
	}
}


BOOST_AUTO_TEST_CASE_TEMPLATE( c_matrixmult, T, complex_types )
{
	T a[9] = {T(1,2),T(3,4),T(5,6), T(7,8),T(9,10),T(11,12), T(13,14),T(15,16),T(17,18)};
	T b[9] = {T(19,20),T(21,22),T(23,24), T(25,26),T(27,28),T(29,30), T(31,32),T(33,34),T(35,36)};
	T c[9] = {T(1,2),T(3,4),T(5,6), T(7,8),T(9,10),T(11,12), T(13,14),T(15,16),T(17,18)};

	T ans[][9] = {
		{T(-87,582),T(-93,624),T(-99,666), T(-105,1500),T(-111,1614),T(-117,1728), T(-123,2418),T(-129,2604),T(-135,2790)},         // a*b
		{T(-99,1290),T(-105,1380),T(-111,1470), T(-105,1596),T(-111,1710),T(-117,1824), T(-111,1902),T(-117,2040),T(-123,2178)},    // a.'*b
		{T(-75,466),T(-93,592),T(-111,718), T(-93,1240),T(-111,1582),T(-129,1924), T(-111,2014),T(-129,2572),T(-147,3130)},         // a*b.'
		{T(-87,+1014),T(-105,+1284),T(-123,+1554),T(-93,+1272),T(-111,+1614),T(-129,+1956),T(-99,+1530),T(-117,+1944),T(-135,+2358)} // a.'*b.'
		};

	warp::arraynd<T, 2> A(warp::extents[3][3], a);
	warp::arraynd<T, 2> B(warp::extents[3][3], b);
	warp::arraynd<T, 2> C(warp::extents[3][3], c);

	// A*B
	warp::la::matrixmult(A, B, C);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[0],ans[0]+9);

	// A.'*B
	warp::la::matrixmult(A, B, C, warp::blas::transpose, warp::blas::dont_transpose);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[1],ans[1]+9);
	
	// A*B.'
	warp::la::matrixmult(A, B, C, warp::blas::dont_transpose, warp::blas::transpose);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[2],ans[2]+9);
	
	// A.'*B.'
	warp::la::matrixmult(A, B, C, warp::blas::transpose, warp::blas::transpose);
	BOOST_CHECK_EQUAL_COLLECTIONS(c,c+9,ans[3],ans[3]+9);
}


BOOST_AUTO_TEST_CASE_TEMPLATE( transpose, T, all_types )
{
	T a[9]    = {1,2,3,       4,5,6,       7,8,9};
	T at[9]   = {1,4,7,       2,5,8,       3,6,9};
	T ans[9]  = {0,0,0,       0,0,0,       0,0,0};

	warp::arraynd<T, 2> A(warp::extents[3][3], a);
	warp::arraynd<T, 2> ANS(warp::extents[3][3], ans);

	warp::la::transpose(A, ANS);

	BOOST_CHECK_EQUAL_COLLECTIONS(at,at+9,ans,ans+9);
}

